#include "pub_sub_common.h"

double total_latency = 0.0;
ipc::route* p_shm_channel = nullptr;
std::chrono::high_resolution_clock::time_point start_time;
std::chrono::high_resolution_clock::time_point end_time;

// Signal handler function
void signal_handler(int signal) {
  fprintf(stdout, "Received signal %d\n", signal);
  if (p_shm_channel != nullptr) {
    p_shm_channel->disconnect();
    fprintf(stdout, "Disconnected from shared memory channel\n");
  }
}

// Function to handle received shared memory data
void shm_consumer(ipc::route& shm_channel) {
  auto buffer = shm_channel.recv();
  const uint32_t buffer_size = buffer.size();

  if (buffer_size < sizeof(uint32_t)) {
    fprintf(stderr, "buffer_size < 4, received termination signal\n");
    return;
  }
  end_time = std::chrono::high_resolution_clock::now();

#if ENABLE_CRC_CHECK
  uint8_t* buffer_header = reinterpret_cast<uint8_t*>(buffer.data());
  uint32_t received_length;
  std::memcpy(&received_length, buffer_header, sizeof(received_length));

  uint16_t received_crc;
  std::memcpy(&received_crc, buffer_header + sizeof(received_length), sizeof(received_crc));

  uint32_t data_length = buffer_size - sizeof(uint32_t) - sizeof(uint16_t);
#if DEBUG_ON
  fprintf(stdout, "buffer_size: %u, received_length: %u, data_length: %u, received_crc: %hx\n", buffer_size,
          received_length, data_length, received_crc);
#endif

  if (received_length != data_length) {
#if DEBUG_ON
    fprintf(stdout, "data length exceeds buffer size, expected %u, got %u\n", received_length, data_length);
#endif
    return;
  }

  uint16_t calculated_crc =
    calculate_crc2(reinterpret_cast<char*>(buffer_header + sizeof(uint32_t) + sizeof(uint16_t)), received_length);

  if (calculated_crc != received_crc) {
    fprintf(stdout, "CRC mismatch: expected %hx, got %hx\n", received_crc, calculated_crc);
    return;
  }
#if DEBUG_ON
  fprintf(stdout, "received data of size %u bytes with valid CRC\n", received_length);
#endif
#endif
}

// Function to run the nanomsg subscriber
void nanomsg_subscriber(std::vector<std::chrono::duration<double, std::micro>>& latencies) {
  ipc::route shm_channel(kIpcRouteName, ipc::receiver);
  if (!shm_channel.valid()) {
    fprintf(stderr, "channel %s is not valid.\n", kIpcRouteName);
    return;
  } else {
    p_shm_channel = &shm_channel;
  }

  setup_signal_handlers();

  int sock = nn_socket(AF_SP, NN_SUB);
  if (sock < 0) {
    fprintf(stderr, "Error creating subscriber socket: %s\n", nn_strerror(nn_errno()));
    return;
  }
  if (nn_connect(sock, URL) < 0) {
    fprintf(stderr, "Error connecting subscriber socket: %s\n", nn_strerror(nn_errno()));
    nn_close(sock);
    return;
  }
  nn_setsockopt(sock, NN_SUB, NN_SUB_SUBSCRIBE, "", 0);
  nn_setsockopt(sock, NN_SOL_SOCKET, NN_RCVTIMEO, &RECEIVE_TIMEOUT, sizeof(RECEIVE_TIMEOUT));

  fprintf(stdout, "Subscriber socket connected to %s\n", URL);

  char* buf = nullptr;
  nn_pollfd pfd[1];
  pfd[0].fd = sock;
  pfd[0].events = NN_POLLIN;
  uint32_t recv_count{0};
  fprintf(stdout, "@@ Start to recv, current time: %s\n", get_currenttime_millis().c_str());
  start_time = std::chrono::high_resolution_clock::now();
  while (recv_count < NUM_MESSAGES) {
    int rc = nn_poll(pfd, 1, POLL_TIMEOUT);
    if (rc < 0) {
      fprintf(stderr, "Error polling socket: %s\n", nn_strerror(nn_errno()));
      return;
    } else if (rc > 0) {
      if ((pfd[0].revents & NN_POLLIN) != 0U) {
        int bytes = nn_recv(sock, &buf, NN_MSG, 0);
        shm_consumer(shm_channel);
        if (bytes < 0) {
          fprintf(stderr, "Error receiving message %d: %s\n", recv_count + 1, nn_strerror(nn_errno()));
          if (buf) {
            nn_freemsg(buf);
          }
          return;
        }
        nn_freemsg(buf);
        recv_count += 1;
      }
    } else {
#if DEBUG_ON
      fprintf(stdout, "nn_poll timeout %d(ms)\n", POLL_TIMEOUT);
#endif
    }
  }
  end_time = std::chrono::high_resolution_clock::now();
  fprintf(stdout, "@@ Finish recv, current time: %s\n", get_currenttime_millis().c_str());
  nn_close(sock);
}

int main() {
  std::vector<std::chrono::duration<double, std::micro>> latencies;
  nanomsg_subscriber(latencies);

  auto total_latency = end_time - start_time;
  auto avg_latency = total_latency / NUM_MESSAGES;
  fprintf(stdout, "Subscriber receive %d messages.\nAverage latency: %lf milliseconds\n", NUM_MESSAGES,
          std::chrono::duration<double, std::milli>(avg_latency).count());

  return 0;
}
